// "Copyright [2021] <Copyright Shengkai Lin>"

#ifndef TASK_ALLREDUCE_H_
#define TASK_ALLREDUCE_H_

#include "ai_sim.h"
#include "network_controller/topology.h"
#include "network_controller/topo_manager.h"

namespace ai_sim {

class AllreduceAlgorithm {
 public:
    virtual void generateCommPairs(taskboard&, TopoManager *) = 0;
    virtual void generateFLows(taskboard&, TopoManager *) = 0;
    static AllreduceAlgorithm* get(TaskType type);
    virtual ~AllreduceAlgorithm(){}
};

class RingAllreduceAlgorithm : public AllreduceAlgorithm {
 public:
    void generateCommPairs(taskboard&, TopoManager * topos);
    void generateFLows(taskboard&, TopoManager * topos);
};

class ButterflyAlgorithm : public AllreduceAlgorithm {
 public:
    void generateCommPairs(taskboard&, TopoManager * topos);
    void generateFLows(taskboard&, TopoManager * topos);
};

}  // namespace ai_sim
#endif  // TASK_ALLREDUCE_H_
